package all_router

import (
	"bytes"
	"chaowan/data_source"
	"chaowan/models"
	"encoding/json"
	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
	"io/ioutil"
	"net/http"
	"sync"
	"time"
)

type UserinfoBody struct {
	Uid int64 `json:"uid" gorm:"uid"`
}

func HelloMiddleware(rl *RateLimiter) gin.HandlerFunc {
	return func(c *gin.Context) {
		body, _ := ioutil.ReadAll(c.Request.Body)
		c.Request.Body = ioutil.NopCloser(bytes.NewBuffer(body))
		var params struct {
			Uid int64 `json:"uid" gorm:"uid"`
		}
		json.Unmarshal(body, &params)
		clientIP := c.ClientIP()
		limiter := rl.GetLimiter(clientIP)
		if !limiter.Allow() {
			c.JSON(http.StatusOK, gin.H{
				"code":    0,
				"error":   0,
				"success": 200,
				"data":    "访问过快",
				"msg":     "ok",
			})
			c.Abort()
			return
		}
		uid := params.Uid
		if uid > 0 {
			rl.RecordUser(uid)
		}

	}
}

// RateLimiter 结构体用于存储每个IP地址的限流器
type RateLimiter struct {
	ipToLimiter   map[string]*rate.Limiter
	mu            sync.Mutex
	recordedUsers map[int64]bool
}

// NewRateLimiter 创建一个新的 RateLimiter 实例
func NewRateLimiter(rps float64, burst int) *RateLimiter {
	return &RateLimiter{
		ipToLimiter:   make(map[string]*rate.Limiter),
		recordedUsers: make(map[int64]bool),
	}
}

// GetLimiter 获取指定IP地址的限流器，如果不存在则创建一个新的
func (rl *RateLimiter) GetLimiter(ip string) *rate.Limiter {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	limiter, exists := rl.ipToLimiter[ip]
	if !exists {
		limiter = rate.NewLimiter(rate.Every(1*time.Second), 5)
		rl.ipToLimiter[ip] = limiter
	}

	return limiter
}

func (ur *RateLimiter) RecordUser(userID int64) bool {

	ur.mu.Lock()
	defer ur.mu.Unlock()
	limiter, exists := ur.recordedUsers[userID]

	if !exists {
		ur.recordedUsers[userID] = true
		data_source.Db.Model(&models.User{}).Where("id = ?", userID).Update("update_time", time.Now().Unix())
	}
	return limiter
}
